import re
from pathlib import Path
from collections import defaultdict

TRACE_FUNCTION = {'$$$$$', '@@@@@@'}

# C/C++ keywords
C_KEYWORDS = {
    'auto', 'break', 'case', 'char', 'const', 'continue', 'default', 'do',
    'double', 'else', 'enum', 'extern', 'float', 'for', 'goto', 'if',
    'int', 'long', 'register', 'return', 'short', 'signed', 'sizeof', 'static',
    'struct', 'switch', 'typedef', 'union', 'unsigned', 'void', 'volatile', 'while',
    '_Alignas', '_Alignof', '_Atomic', '_Bool', '_Complex', '_Generic', '_Imaginary',
    '_Noreturn', '_Static_assert', '_Thread_local'
}

# C++ additional keywords
CPP_KEYWORDS = {
    'alignas', 'alignof', 'and', 'and_eq', 'asm', 'atomic_cancel', 'atomic_commit', 
    'atomic_noexcept', 'auto', 'bitand', 'bitor', 'bool', 'break', 'case', 'catch', 
    'char', 'char8_t', 'char16_t', 'char32_t', 'class', 'compl', 'concept', 'const', 
    'consteval', 'constexpr', 'constinit', 'const_cast', 'continue', 'co_await', 'co_return', 
    'co_yield', 'decltype', 'default', 'delete', 'do', 'double', 'dynamic_cast', 'else', 
    'enum', 'explicit', 'export', 'extern', 'false', 'float', 'for', 'friend', 'goto', 'if', 
    'inline', 'int', 'long', 'mutable', 'namespace', 'new', 'noexcept', 'not', 'not_eq', 'nullptr', 
    'operator', 'or', 'or_eq', 'private', 'protected', 'public', 'reflexpr', 'register', 'reinterpret_cast', 
    'requires', 'return', 'short', 'signed', 'sizeof', 'static', 'static_assert', 'static_cast', 'struct', 
    'switch', 'synchronized', 'template', 'this', 'thread_local', 'throw', 'true', 'try', 'typedef', 'typeid', 
    'typename', 'union', 'unsigned', 'using', 'virtual', 'void', 'volatile', 'wchar_t', 'while', 'xor', 'xor_eq'
}

# Combine C and C++ keywords
ALL_KEYWORDS = C_KEYWORDS.union(CPP_KEYWORDS)

# Control flow statements that take parentheses
CONTROL_FLOW_STATEMENTS = {
    'if', 'for', 'while', 'switch', 'catch'
}

# C++ cast operators
CPP_CAST_OPERATORS = {
    'static_cast', 'dynamic_cast', 'reinterpret_cast', 'const_cast'
}

def find_functions_in_code(source_code: str, is_cpp: bool) -> list[tuple[str, str, int, str]]:
    """Find all function definitions in C/C++ code
    
    Args:
        source_code (str): Source code
        is_cpp (bool): True for C++, False for C
        
    Returns:
        list[tuple[str, str, int, str]]: List of functions (name, return_type, line_number, body)
    """
    code = remove_comments(source_code)
    
    # Improved pattern to match function signature more accurately
    if is_cpp:
        signature_pattern = re.compile(r'''
            (?:template\s*<[^>]*>\s*)?  # Optional template declaration
            ([\w\:\*&\s]+?)            # Return type
            \s*
            ([\w\:]+)                  # Function name
            \s*
            \(([^)]*)\)                # Parameter list
            \s*
            (?:const)?\s*              # Optional const qualifier
            \{                         # Function body start
        ''', re.VERBOSE)
    else:
        signature_pattern = re.compile(r'''
            (\w+\s*[\*\s]*?)  # Return type
            \s+
            (\w+)             # Function name
            \s*
            \(([^)]*)\)       # Parameter list
            \s*
            \{                # Function body start
        ''', re.VERBOSE)
    
    functions = []
    
    for match in signature_pattern.finditer(code):
        return_type = match.group(1).strip()
        function_name = match.group(2).strip()
        start_pos = match.end() - 1  # Position of the opening brace {
        
        # Find the matching closing brace using a stack
        body, end_pos = extract_balanced_block(code, start_pos)
        if body is None:
            continue
            
        line_number = code[:start_pos].count('\n')
        
        if function_name not in ALL_KEYWORDS:
            functions.append((function_name, return_type, line_number, body))
    
    return functions

def extract_balanced_block(code: str, start_pos: int) -> tuple[str, int]:
    """Extract a balanced block of code starting at start_pos, handling nested braces"""
    if start_pos >= len(code) or code[start_pos] != '{':
        return None, start_pos
        
    stack = ['{']
    end_pos = start_pos + 1
    
    while end_pos < len(code):
        char = code[end_pos]
        if char == '{':
            stack.append('{')
        elif char == '}':
            stack.pop()
            if not stack:
                # Found matching closing brace
                return code[start_pos+1:end_pos], end_pos + 1
                
        end_pos += 1
    
    # If we reach here, there's an unmatched opening brace
    return code[start_pos+1:], end_pos

def extract_function_calls(function_name: str, function_body: str, is_cpp: bool, file_path: Path, function_line: int) -> list[tuple[str, Path, int]]:
    """Extract function calls from a function body
    
    Args:
        function_body (str): Function body code
        is_cpp (bool): True for C++, False for C
        file_path (Path): Path to the source file
        
    Returns:
        list[tuple[str, Path, int]]: List of called function names, file paths, and line numbers
    """
    # Remove strings and character literals to avoid false matches
    code = re.sub(r'"[^"]*"', '', function_body)
    code = re.sub(r"'[^']*'", '', code)
    
    # Remove comments
    code = remove_comments(code)
    
    # Split code into statements to better identify control flow
    statements = split_into_statements(code)
    
    if function_name in TRACE_FUNCTION:
        print(f'{function_name}{statements}')

    calls = []
    for stmt in statements:
        stmt = stmt.strip()
        if not stmt:
            continue
            
        # Basic function call pattern
        if is_cpp:
            call_pattern = re.compile(r'''
                ([\w:]+)            # Function name (allow namespaces)
                \s*
                \(([^)]*)\)         # Argument list
                (?:\s*->\s*[\w:]+)? # Optional member access after call
            ''', re.VERBOSE)
        else:
            call_pattern = re.compile(r'''
                (\w+)               # Function name
                \s*
                \(([^)]*)\)         # Argument list
            ''', re.VERBOSE)
        
        if function_name in TRACE_FUNCTION:
            print(f'{stmt}')

        for match in call_pattern.finditer(stmt):
            func_name = match.group(1)
            
            # Split namespaced name
            components = func_name.split('::')
            last_component = components[-1]
            
            if function_name in TRACE_FUNCTION:
                print(f'{func_name}, {match.group(2)}, {last_component}')

            # Filter out keywords
            if is_cpp:
                if all(comp not in ALL_KEYWORDS for comp in components):
                    # Calculate line number
                    pos = function_body.find(stmt) + match.start()
                    line_number = function_body[:pos].count('\n') + 1
                    calls.append((func_name, file_path, line_number + function_line))
            else:
                if func_name not in C_KEYWORDS:
                    # Calculate line number
                    pos = function_body.find(stmt) + match.start()
                    line_number = function_body[:pos].count('\n') + 1
                    calls.append((func_name, file_path, line_number + function_line))
            
            # Argument list have function calls
            arg_list = match.group(2)
            arg_pattern = re.compile(r'([\w:]+)\s*\(', re.VERBOSE)

            if function_name in TRACE_FUNCTION:
                print(f'arg_list: {arg_list}')

            for arg_match in arg_pattern.finditer(arg_list):
                arg_func_name = arg_match.group(1)
                if arg_func_name not in ALL_KEYWORDS and arg_func_name not in CONTROL_FLOW_STATEMENTS:
                    if function_name in TRACE_FUNCTION:
                        print(f'{function_name} arg-> {arg_func_name}')
                    # Calculate line number
                    arg_pos = function_body.find(stmt) + match.start() + arg_match.start()
                    line_number = function_body[:arg_pos].count('\n') + 1
                    calls.append((arg_func_name, file_path, line_number + function_line))
    
    seen = set()
    unique_calls = []
    for call in calls:
        if call not in seen:
            seen.add(call)
            unique_calls.append(call)
    
    return unique_calls

def split_into_statements(code: str) -> list[str]:
    """Split code into individual statements
    
    This is a simplified approach that handles basic cases.
    Does not handle nested blocks perfectly but works for most common cases.
    """
    statements = []
    current = ""
    paren_count = 0
    brace_count = 0
    bracket_count = 0
    
    for char in code:
        current += char
        
        # Track nested parentheses, braces, and brackets
        if char == '(':
            paren_count += 1
        elif char == ')':
            paren_count -= 1
        elif char == '{':
            brace_count += 1
        elif char == '}':
            brace_count -= 1
        elif char == '[':
            bracket_count += 1
        elif char == ']':
            bracket_count -= 1
            
        # End of statement when all brackets are closed and a semicolon is encountered
        if char == ';' and paren_count == 0 and brace_count == 0 and bracket_count == 0:
            statements.append(current.strip())
            current = ""
    
    # Add any remaining code (last statement without trailing semicolon)
    if current.strip():
        statements.append(current.strip())
    
    return statements

def remove_comments(source_code: str) -> str:
    """Remove comments from C/C++ code"""
    # Remove line comments
    code = re.sub(r'//.*?$', '', source_code, flags=re.MULTILINE)
    # Remove block comments
    code = re.sub(r'/\*.*?\*/', '', code, flags=re.DOTALL)
    return code

def process_source_file(file_path: Path) -> list[tuple[str, str, int, Path, list[tuple[str, Path, int]]]]:
    """Process a source file and return functions with call information
    
    Returns:
        list[tuple[str, str, int, Path, list[tuple[str, Path, int]]]]: 
            List of functions (name, return_type, line_number, file_path, called_functions)
    """
    try:
        with open(file_path, 'r', encoding='utf-8') as file:
            content = file.read()
        
        is_cpp = file_path.suffix.lower() in {'.cpp', '.cxx', '.cc', '.c++', '.hpp', '.hxx', '.hh', '.h++'}
        
        functions = find_functions_in_code(content, is_cpp)
        result = []
        
        for name, ret_type, line, body in functions:
            called_functions = extract_function_calls(name, body, is_cpp, file_path, line)
            result.append((name, ret_type, line, file_path, called_functions))
        
        return result
    except UnicodeDecodeError:
        print(f"Error: File '{file_path}' is not UTF-8 encoded")
        return []
    except Exception as e:
        print(f"Error processing file '{file_path}': {e}")
        return []

def find_source_files(directory: Path) -> list[Path]:
    """Find all C/C++ source files in a directory recursively"""
    c_extensions = {'*.c', '*.h'}
    cpp_extensions = {'*.cpp', '*.cxx', '*.cc', '*.c++', '.hpp', '*.hxx', '*.hh', '*.h++'}
    
    all_files = []
    for ext in c_extensions.union(cpp_extensions):
        all_files.extend(directory.rglob(ext))
    
    return all_files

def find_callers(target_function, all_functions):
    callers = []
    for name, _, _, _, called_functions in all_functions:
        for called, _, _ in called_functions:
            if called == target_function:
                callers.append(name)
    return callers

def find_reffers(target_function, all_functions):
    for name, _, _, _, called_functions in all_functions:
        if target_function in name:
            return called_functions
    return []

def build_call_stacks_top_down(target_function, all_functions, current_stack=None):
    if current_stack is None:
        current_stack = [target_function]
    for name, _, _, _, called_functions in all_functions:
        if name == target_function:
            for called, _, _ in called_functions:
                if called in current_stack:
                    continue  # 避免循环调用
                new_stack = current_stack + [called]
                reffers = find_reffers(called, all_functions)
                if not reffers:
                    yield new_stack
                yield from build_call_stacks_top_down(called, all_functions, new_stack)


def build_call_stacks_bottom_up(target_function, all_functions, current_stack=None):
    if current_stack is None:
        current_stack = [target_function]
    callers = find_callers(target_function, all_functions)
    if not callers:
        yield current_stack
    for caller in callers:
        if caller in current_stack:
            continue  # 避免循环调用
        new_stack = [caller] + current_stack
        yield from build_call_stacks_bottom_up(caller, all_functions, new_stack)
def main():
    """Main function for command-line usage"""
    import argparse

    parser = argparse.ArgumentParser(description='Analyze C/C++ functions and their calls')
    parser.add_argument('path', help='C/C++ source file or directory path')
    parser.add_argument('--calls', action='store_true', help='Show functions called by each function')
    parser.add_argument('--call-stacks-top', type=str, help='Show all call stacks with a specific function as the top')
    parser.add_argument('--call-stacks-bottom', type=str, help='Show all call stacks with a specific function as the bottom')
    args = parser.parse_args()

    path = Path(args.path)

    if not path.exists():
        print(f"Error: Path '{path}' does not exist")
        return

    all_functions = []

    if path.is_file():
        if path.suffix.lower() not in {'.c', '.h', '.cpp', '.cxx', '.cc', '.c++', '.hpp', '.hxx', '.hh', '.h++'}:
            print(f"Warning: File '{path}' may not be a C/C++ source file")
        all_functions = process_source_file(path)
    else:
        source_files = find_source_files(path)
        if not source_files:
            print(f"Error: No C/C++ files found in directory '{path}'")
            return

        print(f"Found {len(source_files)} source files in directory '{path}', analyzing...")

        for file in source_files:
            file_functions = process_source_file(file)
            all_functions.extend(file_functions)

    if not all_functions:
        print("No function definitions found")
        return

    if args.calls:
        print(f"Total {len(all_functions)} functions found with call information:")
        print("=" * 120)
        print(f"{'Function Name':<30} {'Called Functions'}")
        print("-" * 120)

        for name, ret_type, line, file_path, called in all_functions:
            # 组合文件路径和行号
            file_path_with_line = f"{file_path}:{line}"
            called_str = ', '.join([f"{func} ({path}:{ln})" for func, path, ln in called]) if called else '-'
            print(f"{name:<30} {called_str}")
    elif args.call_stacks_top:
        print(f"All complete call stacks with function '{args.call_stacks_top}' as the top:")
        for stack in build_call_stacks_top_down(args.call_stacks_top, all_functions):
            print(" -> ".join(stack))
    elif args.call_stacks_bottom:
        print(f"All complete call stacks with function '{args.call_stacks_bottom}' as the bottom:")
        for stack in build_call_stacks_bottom_up(args.call_stacks_bottom, all_functions):
            print(" -> ".join(stack))
    else:
        print(f"Total {len(all_functions)} functions found:")
        print("=" * 60)
        print(f"{'Function Name':<30} {'File Path:Line'}")
        print("-" * 60)
        for name, _, line, file_path, _ in all_functions:
            file_path_with_line = f"{file_path}:{line}"
            print(f"{name:<30} {file_path_with_line}")


if __name__ == "__main__":
    main()